# coding = utf-8

'''
save the data (only contains the liver)
'''

import os
import numpy as np
import cv2
from scipy import ndimage
from skimage import measure
import prettytable as pt
from pathlib2 import Path
import math

def save_only_liver():
    data_path = "/datasets/3DIRCADB/chengkun_remove"
    save_path = "/datasets/3DIRCADB/chengkun_remove_only_liver"

    for i in range(20):
        case_id = "case_{}".format(str(i).zfill(5))
        case_path = os.path.join(data_path, case_id)
        segmentation_path = os.path.join(case_path, "segmentation")
        image_path = os.path.join(case_path, "imaging")

        destination_patient_path = os.path.join(save_path, "case_{}".format(str(i).zfill(5)))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        j = 0
        for item in sorted(os.listdir(segmentation_path)):
            segmentation_file = os.path.join(segmentation_path, item)
            segmentation = np.load(segmentation_file)
            image_file = os.path.join(image_path, item)
            image = np.load(image_file)
            if np.max(segmentation) <= 0:
                continue

            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(j).zfill(3))), image)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(j).zfill(3))), segmentation)
            j += 1

def save_only_tumor():
    data_path = "/datasets/3DIRCADB/chengkun_remove"
    save_path = "/datasets/3DIRCADB/chengkun_remove_only_tumor"

    for i in range(20):
        case_id = "case_{}".format(str(i).zfill(5))
        case_path = os.path.join(data_path, case_id)
        segmentation_path = os.path.join(case_path, "segmentation")
        image_path = os.path.join(case_path, "imaging")

        destination_patient_path = os.path.join(save_path, "case_{}".format(str(i).zfill(5)))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        j = 0
        for item in sorted(os.listdir(segmentation_path)):
            segmentation_file = os.path.join(segmentation_path, item)
            segmentation = np.load(segmentation_file)
            image_file = os.path.join(image_path, item)
            image = np.load(image_file)
            if np.max(segmentation) <= 1:
                continue

            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(j).zfill(3))), image)
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(j).zfill(3))), segmentation)
            j += 1

#对肿瘤数据进行分析
def analysis_only_tumor():
    data_path = "/datasets/3DIRCADB/chengkun_remove_only_tumor"
    max_tumor = 0
    tumor_size = 0
    background_size = 0
    liver_size = 0

    for i in range(20):
        case_id = "case_{}".format(str(i).zfill(5))
        case_path = os.path.join(data_path, case_id)
        segmentation_path = os.path.join(case_path, "segmentation")
        image_path = os.path.join(case_path, "imaging")
        np_list = []

        for item in sorted(os.listdir(segmentation_path)):
            item_file = os.path.join(segmentation_path, item)
            data = np.load(item_file)
            np_list.append(data)

        data = np.array(np_list)

        tumor_size += (data >= 2).sum()
        background_size += (data == 0).sum()
        liver_size += (data == 1).sum()

        if data is None or len(data) == 0:
            continue

        desp = []
        for j in range(2, np.max(data)+1):
            desp.append("{}:{}".format(j, (data == j).sum()))
        print("{}\t{}\t".format(i, ",".join(desp)), (data >= 2).sum())

    print(max_tumor, background_size, liver_size, tumor_size)


#针对数据进行比对
def analysis_tumor():
    tumor = []
    with open("../only_tumor_analysis", "r") as file:
        for line in file:
            data = line.strip().split("\t")[1]
            for item in data.split(","):
                tumor.append(int(item.split(":")[1]))

    size = len(tumor)
    print(size)
    tumor = np.array(tumor)
    avg_tumor = np.sum(tumor)/len(tumor)
    print(avg_tumor, np.max(tumor))

    big = 0
    middle = 0
    small = 0
    for item in tumor:
        if item >= 100000:
            big += 1
        elif item >= 10000:
            middle += 1
        else:
            small += 1

    print(big, big/size)
    print(middle, middle/size)
    print(small, small/size)


if __name__ == '__main__':
    analysis_tumor()

